import os
import Core.Config as Config
import Core.Gadget as Gadget
import pandas as pd
import datetime
import numpy as np


def get_etf_return_data(database, datetime1, datetime2, symbol_list=[], test=False):
    return get_return_data(database, datetime1, datetime2, symbol_list, test, instrument_type="etf")

def get_index_return_data(database, datetime1, datetime2, symbol_list=[], test=False):
    return get_return_data(database, datetime1, datetime2, symbol_list, test, instrument_type="index")


def get_return_data(database, datetime1, datetime2, symbol_list=[], test=False, instrument_type="index"):

    #
    df_index_list = database.GetDataFrame("financial_data", "instrument_" + instrument_type)

    if test:
        df_index_list = df_index_list[:100]

    #
    max_count = 0
    df_bars = pd.DataFrame()
    for index, row in df_index_list.iterrows():
        symbol = row["symbol"]
        name = row["description"]
        # print(index, symbol, name)
        # if index == 258:
        #     aa = 0

        if instrument_type == "etf" and row["invest_type1"] == "货币市场型基金":
            continue

        # index_dailybar
        #
        df_index = database.GetDataFrame("financial_data", instrument_type + "_dailybar",
                                         filter=[("symbol", symbol), ("date",">=", datetime1), ("date", "<=", datetime2)],
                                         sort=[("date", 1)])
        if df_index.empty:
            continue

        # 所有值都相同
        if df_index['close'].nunique() == 1:
            continue

        # 以第一个得到的数据序列长度为标准长度
        if max_count == 0:
            max_count = len(df_index)

        # df_index["daily_return"] = df_index["close"] / df_index["close"].shift(1) - 1
        df_bars_tmp = df_index[["date", "close"]].copy()
        df_bars_tmp.dropna(subset=["close"],inplace=True) # 不计空数据

        # 数据缺失太多，放弃
        if len(df_bars_tmp) / max_count < 0.5:
            continue

        df_bars_tmp.rename(columns={"close":symbol}, inplace=True)
        #
        if df_bars.empty:
            df_bars = df_bars_tmp
        else:
            df_bars = pd.merge(df_bars, df_bars_tmp, how="outer", on="date")

    df_bars.set_index("date", inplace=True)
    df_bars.sort_values(by="date", ascending=True, inplace=True) # outer, 必须重新排序
    df_bars.ffill(inplace=True)  # 不同asset交易不同，需要填补空白
    df_returns = np.log(df_bars / df_bars.shift(1)) # 转换成对数收益率
    df_returns = df_returns.iloc[1:].copy()  # 抛弃第一行
    return df_returns



def test_asset_corelation(database):
    bond_symbol = "CBA00301.CS" # 中债-总财富(总值)指数
    df_bond = database.GetDataFrame("financial_data", "index_dailybar", filter=[("symbol",bond_symbol)], sort=[("date",1)])
    df_bond["daily_return"] = df_bond["close"] / df_bond["close"].shift(1) - 1

    datetime1 = datetime.datetime(2022,1,1)
    datetime2 = datetime.datetime(2024,12,27)
    df_bond = df_bond[(df_bond["date"]>=datetime1) & (df_bond["date"]<=datetime2)].copy()

    df_index_list = database.GetDataFrame("financial_data", "instrument_index")
    for index, row in df_index_list.iterrows():
        test_symbol = row["symbol"]
        test_name = row["description"]
        df_index = database.GetDataFrame("financial_data", "index_dailybar", filter=[("symbol", test_symbol)], sort=[("date", 1)])
        df_index["daily_return"] = df_index["close"] / df_index["close"].shift(1) - 1

        df = pd.merge(df_bond[["date","daily_return"]], df_index[["date","daily_return"]], how="inner", on="date")
        df.dropna(inplace=True)

        # 计算相关性
        pearson_corr = df["daily_return_x"].corr(df["daily_return_y"], method="pearson")  # spearman
        spearman_corr = df["daily_return_x"].corr(df["daily_return_y"], method="spearman")
        #
        print(bond_symbol, test_name, test_symbol, pearson_corr, spearman_corr)
        aa = 0


if __name__ == '__main__':

    # ---Connecting Database---
    path_filename = os.getcwd() + "\..\Config\config_local.json"
    database = Config.create_database(database_type="MySQL", config_file=path_filename, config_field="MySQL")

    # test_asset_corelation(database)

    datetime1 = datetime.datetime(2024,1,1)
    datetime2 = datetime.datetime(2024,12,31)
    get_index_return_data(database, datetime1, datetime2, symbol_list=[])
